// EG4EventAction.hh -- a brief definition for the class
// Created: Jan. 27, 2019, transcribed by Asia SUN, asia.rabbit@163.com
// Last modified: Nov. 12, 2023, by Asia SUN

#include <array>
#include "EG4EventAction.h"
#include "EG4DCHit.h"
#include "EG4PlaHit.h"
#include "EG4RunAction.h"
#include "G4Event.hh"
#include "G4UnitsTable.hh"
#include "G4RunManager.hh"
#include "G4SDManager.hh"
#include "G4HCofThisEvent.hh"
#include "G4SystemOfUnits.hh"
#include "TTree.h"
#include "EG4Data.h"

using std::array;
using std::cout;
using std::endl;

G4VHitsCollection *EG4EventAction::GetHC(const G4Event *ev, int id){
  return ev->GetHCofThisEvent()->GetHC(id);
} // end member function GetHC

EG4EventAction::EG4EventAction(EG4RunAction *runAction) : G4UserEventAction(),
fPlaHCID{0}, fDCHCID{0}, fNWHCID{0}, fRunAction(runAction){
  for(auto &t : fPlaHCID) t = -1;
  for(auto &t : fDCHCID) t = -1;
  for(auto &t : fNWHCID) t = -1;
  fData = fRunAction->data();
} // end ctor

EG4EventAction::~EG4EventAction(){}

void EG4EventAction::BeginOfEventAction(const G4Event *ev){
  static auto sdm = G4SDManager::GetSDMpointer();
  // find hit collections by name (just once)
  // and save them in the data members of this class
  if(-1 == fPlaHCID[0]){ // assign only once
    // hits collections names
    array<G4String, 2> pHCName // plastic scintillator hit collection name
      = {{"T1_SD/PlaC", "TOFWStrip_SD/PlaC"}}; // SDname/HCName  - primary
    array<G4String, 2> dHCName // drift chamber hit collection name
      = {{"DCTaD_SD/DCC", "DC_SD/DCC"}}; //  - primary
    array<G4String, 3> nwHCName // neutron wall hit collection name - secondaries
      = {{"NWSciStrip_SD/PlaC2N", "NWFeStrip_SD/PlaC2N", "NWVetoStrip_SD/PlaC2"}};

    for(int i = 0; i < 2; i++){
      // hit collections IDs
      fPlaHCID[i] = sdm->GetCollectionID(pHCName[i]);
      fDCHCID[i] = sdm->GetCollectionID(dHCName[i]);
    } // end for over i
    for(int i = 0; i < 3; i++) fNWHCID[i] = sdm->GetCollectionID(nwHCName[i]);
  } // end if
  fData->Initialize();
} // end member function BeginOfEventAction

void EG4EventAction::EndOfEventAction(const G4Event *ev){
  fData->index = ev->GetEventID();
  G4PrimaryVertex *vertex = ev->GetPrimaryVertex(0);
  G4PrimaryParticle *primary = vertex->GetPrimary(0);
  fData->a = primary->GetG4code()->GetPDGMass() / CLHEP::amu_c2;
  fData->z = primary->GetCharge();
  EG4Data::copy(fData->pos, vertex->GetPosition()); // the initial position
  EG4Data::copy(fData->p0, primary->GetMomentum()); // the initial momentum

  // hits belonging to the primaries (0-3) and secondaries (4-6) //
  G4VHitsCollection *hc[7] = {GetHC(ev, fPlaHCID[0]), GetHC(ev, fPlaHCID[1]),
    GetHC(ev, fDCHCID[0]), GetHC(ev, fDCHCID[1]), GetHC(ev, fNWHCID[0]),
    GetHC(ev, fNWHCID[1]), GetHC(ev, fNWHCID[2])};
  for(int i = 7; i--;) if(!hc[i])
    G4Exception("EG4EventAction", "EndOfEventAction()", FatalException,
      "The hit collection not found");
  if((fData->nhT1 = hc[0]->GetSize())){
    (static_cast<EG4PlaHit *>(hc[0]->GetHit(0)))->Assign(fData->hT1);
  } // end if
  if((fData->nhTOFW = hc[1]->GetSize())){
    (static_cast<EG4PlaHit *>(hc[1]->GetHit(0)))->Assign(fData->hTOFW[0]);
  } // end if
  if(fData->nhTOFW >= 2){
    (static_cast<EG4PlaHit *>(hc[1]->GetHit(1)))->Assign(fData->hTOFW[1]);
  } // end if
  fData->nhDCTaD = hc[2]->GetSize();
  for(int i = 0; i < fData->nhDCTaD; i++){ // DCTaD
    EG4DCHit *hit = static_cast<EG4DCHit *>(hc[2]->GetHit(i));
    const int id = hit->GetDCId()*4 + hit->GetWireLayerId();
    hit->Assign(fData->hDCTaD[id]);
  } // end for over i and j
  fData->nhDC = hc[3]->GetSize();
  for(int i = 0; i < fData->nhDC; i++){ // DC
    EG4DCHit *hit = static_cast<EG4DCHit *>(hc[3]->GetHit(i));
    const int id = hit->GetDCId()*6 + hit->GetWireLayerId();
    hit->Assign(fData->hDC[id]);
  } // end for over i and j
  fData->nhNWS = hc[4]->GetSize();
  for(int i = 0; i < fData->nhNWS; i++){ // sci-strips
    EG4PlaHit *hit = static_cast<EG4PlaHit *>(hc[4]->GetHit(i));
    hit->Assign(fData->hNWS[hit->GetLayerId()]);
  } // end for over hits in DCs
  fData->nhNWF = hc[5]->GetSize();
  for(int i = 0; i < fData->nhNWF; i++){ // Fe-strips
    EG4PlaHit *hit = static_cast<EG4PlaHit *>(hc[5]->GetHit(i));
    hit->Assign(fData->hNWF[hit->GetLayerId()]);
  } // end for over hits in DCs
  if((fData->nhNWV = hc[6]->GetSize())){
    (static_cast<EG4PlaHit *>(hc[6]->GetHit(0)))->Assign(fData->hNWV);
  } // end if

  fData->PostAssign();
  fRunAction->Fill();

  // print diagnostics //
  if(!(ev->GetEventID() % 100)){
    cout << ">>> Event " << ev->GetEventID() << " >>> Simulation truth: "
      << primary->GetG4code()->GetParticleName() << endl; // "\r" << std::flush;
    cout << "Initial P: (" << fData->p0[0] << ", " << fData->p0[1] << ", ";
    cout << fData->p0[2] << ")=" << fData->pp0 << "MeV/c, ";
    cout << primary->GetKineticEnergy() / fData->a << "MeV/u" << endl;
    cout << "Initial Pos: (" << fData->pos[0] << ", " << fData->pos[1] << ", ";
    cout << fData->pos[2] << ")" << endl;
    cout << "Theta: " << fData->theta << " degree, Phi: ";
    cout << fData->phi << " degree" << endl;

    cout << endl << "------ T1 has " << fData->nhT1 << " hits --------" << endl;
    // for(int i = 0; i < fData->nhT1; i++) hc[0]->GetHit(i)->Print();
    // cout << endl << "------ DCTaD has " << fData->nhDCTaD << " hits --------" << endl;
    // for(int i = 0; i < fData->nhDCTaD; i++) hc[2]->GetHit(i)->Print();
    // cout << endl << "------ DC has " << fData->nhDC << " hits --------" << endl;
    // for(int i = 0; i < fData->nhDC; i++) hc[3]->GetHit(i)->Print();
    cout << endl << "------ TOFW has " << fData->nhTOFW << " hits --------" << endl;
    // for(int i = 0; i < fData->nhTOFW; i++) hc[1]->GetHit(i)->Print();
    // print tof2
    if(fData->nhTOFW && fData->nhT1){
      const double t_T1 = fData->hT1.t;
      const double t_TOFW = fData->hTOFW[0].t;
      cout << "\033[36;1mtof2: " << t_TOFW << "-" << t_T1 << "=";
      cout << t_TOFW - t_T1 << " ns\033[0m" << endl;
    } // end if
    // neutron wall
    // cout << endl << "------ NW Sci has " << fData->nhNWS << " hits --------" << endl;
    // for(int i = 0; i < fData->nhNWS; i++) hc[4]->GetHit(i)->Print();
    // cout << endl << "------ NW Fe has " << fData->nhNWF << " hits --------" << endl;
    // for(int i = 0; i < fData->nhNWF; i++) hc[5]->GetHit(i)->Print();
    // cout << endl << "------ NW Veto has " << fData->nhNWV << " hits --------" << endl;
    // for(int i = 0; i < fData->nhNWV; i++) hc[6]->GetHit(i)->Print();
  } // end if
} // end member function EndOfEventAction
